import pymysql
import csv
import argparse


def insert_data_to_mysql(host, user, password, database, csv_file_path):
    # 创建MySQL连接
    connection = pymysql.connect(
        host=host,
        user=user,
        password=password,
        database=database,
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )

    try:
        with connection.cursor() as cursor:
            # 创建表的SQL语句
            create_table_sql = """
            CREATE TABLE IF NOT EXISTS source_file (
                id VARCHAR(255) PRIMARY KEY,
                ABSOLUTEPATH VARCHAR(255),
                SYSFILENAME VARCHAR(255),
                OPERATIONCODE VARCHAR(255),
                OPERATIONID VARCHAR(255),
                PHYFILENAME VARCHAR(255),
                FILETYPE VARCHAR(255),
                FILESIZE INT,
                PERSONID VARCHAR(255),
                PERSONNAME VARCHAR(255),
                OPERATIONNAME VARCHAR(255),
                file_not_found_path VARCHAR(255),
                processed TINYINT(1) DEFAULT 0
            )
            """
            cursor.execute(create_table_sql)

            # 读取CSV文件
            with open(csv_file_path, newline='', encoding='utf-8') as csvfile:
                csv_reader = csv.DictReader(csvfile)

                # 准备批量插入数据
                batch_data = []
                batch_size = 1000

                # 插入数据的SQL语句
                insert_data_sql = """
                INSERT INTO source_file (
                    id, ABSOLUTEPATH, SYSFILENAME, OPERATIONCODE, OPERATIONID, 
                    PHYFILENAME, FILETYPE, FILESIZE, PERSONID, PERSONNAME, 
                    OPERATIONNAME, file_not_found_path, processed
                ) VALUES (
                    %(id)s, %(ABSOLUTEPATH)s, %(SYSFILENAME)s, %(OPERATIONCODE)s, %(OPERATIONID)s, 
                    %(PHYFILENAME)s, %(FILETYPE)s, %(FILESIZE)s, %(PERSONID)s, %(PERSONNAME)s, 
                    %(OPERATIONNAME)s, %(file_not_found_path)s, 0
                )
                """

                # 遍历CSV中的每一行
                for row in csv_reader:
                    batch_data.append(row)
                    # 如果批量数据达到设定大小，执行批量插入
                    if len(batch_data) == batch_size:
                        cursor.executemany(insert_data_sql, batch_data)
                        connection.commit()
                        batch_data = []  # 清空批量数据

                # 插入剩余的不足1000条的数据
                if batch_data:
                    cursor.executemany(insert_data_sql, batch_data)
                    connection.commit()

    finally:
        # 关闭数据库连接
        connection.close()

    print("数据已成功导入到source_file表中。")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='导入CSV数据到MySQL数据库')
    parser.add_argument('--host', type=str, default='localhost', help='MySQL服务器地址')
    parser.add_argument('--user', type=str, required=True, help='MySQL用户名')
    parser.add_argument('--password', type=str, required=True, help='MySQL密码')
    parser.add_argument('--database', type=str, required=True, help='MySQL数据库名称')
    parser.add_argument('--csv_file_path', type=str, required=True, help='CSV文件路径')

    args = parser.parse_args()

    insert_data_to_mysql(
        host=args.host,
        user=args.user,
        password=args.password,
        database=args.database,
        csv_file_path=args.csv_file_path
    )
